{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cartpole Example\n",
    "This notebook goes over the process of setting up and solving a trajectory optimization problem using multiple solvers.\n",
    "\n",
    "## Loading the Required Packages\n",
    "To set up the problem, we import `TrajectoryOptimization`, along with `StaticArrays` and `LinearAlgebra` to help in the set up.\n",
    "\n",
    "To keep things simple, we use the cartpole model from the [`RobotZoo`](https://github.com/bjack205/RobotZoo.jl) package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m\u001b[1m Activating\u001b[22m\u001b[39m environment at `~/.julia/dev/TrajectoryOptimization/examples/Project.toml`\n",
      "┌ Info: Precompiling RobotZoo [74be38bb-dcc2-4b9e-baf3-d6373cd95f10]\n",
      "└ @ Base loading.jl:1260\n"
     ]
    }
   ],
   "source": [
    "import Pkg; Pkg.activate(@__DIR__); Pkg.instantiate();\n",
    "using TrajectoryOptimization\n",
    "using RobotDynamics\n",
    "import RobotZoo.Cartpole\n",
    "using StaticArrays, LinearAlgebra"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the model\n",
    "Here we instantiate our model, which is a sub-type of `AbstractModel`, and extract out the state dimension `n` and control dimension `m`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Cartpole()\n",
    "n,m = size(model);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the model discretization\n",
    "We now define how we will discretize our model time. Here we choose to solve a problem with 5 second horizon and use 101 knotpoints, for a uniform time step of 0.05 seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.05"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "N = 101\n",
    "tf = 5.\n",
    "dt = tf/(N-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set initial and final state\n",
    "All trajectory optimization problems require an initial condition (it will be assumed to be the origin if not passed to the problem constructor). The final state is not be required, and could be replaced with constraints on the terminal state, in practice. Here we want the cartpole to start with the pendulum hanging down and then swing up. Note that we use `SVector`s here, but the constructor will automatically convert normal vectors, as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "x0 = @SVector zeros(n)\n",
    "xf = @SVector [0, pi, 0, 0];  # i.e. swing up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the Objective\n",
    "Here we define a simple LQR tracking objective:\n",
    "$$ \\frac{1}{2} (x_N - x_f)^T Q_f (x_N - x_f) +  \\frac{1}{2} \\sum_{k=0}^{N-1} (x_k - x_f)^T Q_k (x_k - x_f) + (u_k - u_f)^T R_k (u_k - u_f) $$\n",
    "\n",
    "Here it is very important to use the `Diagonal` and `SVector` constructors, since this will have a significant impact on performance. To create the objective, we use the convenient `LQRObjective` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up\n",
    "Q = 1.0e-2*Diagonal(@SVector ones(n))\n",
    "Qf = 100.0*Diagonal(@SVector ones(n))\n",
    "R = 1.0e-1*Diagonal(@SVector ones(m))\n",
    "obj = LQRObjective(Q,R,Qf,xf,N);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Constraints\n",
    "We now add constraints to our problem. We place simple bounds on the controls and make the goal state and hard constraint. Note that we first create an empty `ConstraintList`, then create the individual constraints, and finally add the constraints to the list, specifying the knot point indicies to which the constraint applies. Since we don't have controls at the final knot point, we set the indices to `1:N-1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Empty ConstraintList\n",
    "conSet = ConstraintList(n,m,N)\n",
    "\n",
    "# Control Bounds\n",
    "u_bnd = 3.0\n",
    "bnd = BoundConstraint(n,m, u_min=-u_bnd, u_max=u_bnd)\n",
    "add_constraint!(conSet, bnd, 1:N-1)\n",
    "\n",
    "# Goal Constraint\n",
    "goal = GoalConstraint(xf)\n",
    "add_constraint!(conSet, goal, N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the problem\n",
    "We now pass all the required information to the `Problem` constructor to define our problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "prob = Problem(model, obj, xf, tf, x0=x0, constraints=conSet);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialization\n",
    "As is standard for all non-convex optimization problems, initialization is very important. If a good guess isn't available, it is usually best to initialize the system will a set of controls that keep the state bounded or constant. It is often useful to perturb the system slightly, and provide a dynamicall feasible state trajectory. Here we initialize the system will near-zero controls and then simulate the system forward with the `rollout!` method to provide the initial state guess."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "u0 = @SVector fill(0.01,m)\n",
    "U0 = [u0 for k = 1:N-1]\n",
    "initial_controls!(prob, U0)\n",
    "rollout!(prob);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solving with ALTRO\n",
    "Now that the problem is defined, we can use any of the supported solvers to solve the problem. Here we use [`Altro.jl`](https://github.com/RoboticExplorationLab/ALTRO.jl), a very fast specialized solver that uses iterative LQR (iLQR) within an augmented Lagrangrian framework to handle the constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32;1m\n",
      "SOLVE COMPLETED\n",
      "\u001b[0m solved using the \u001b[0m\u001b[36;1mALTRO\u001b[0m Solver,\n",
      " part of the Altro.jl package developed by the REx Lab at Stanford and Carnegie Mellon Universities\n",
      "\u001b[34;1m\n",
      "  Solve Statistics\n",
      "\u001b[0m    Total Iterations: 40\n",
      "\u001b[0m    Solve Time: 12.698675999999999 (ms)\n",
      "\u001b[34;1m\n",
      "  Covergence\n",
      "\u001b[0m    Terminal Cost: 1.552558743680986\n",
      "\u001b[0m    Terminal dJ: \u001b[31m0.0007412366246843938\n",
      "\u001b[0m    Terminal gradient: \u001b[32m0.003248506450786873\n",
      "\u001b[0m    Terminal constraint violation: \u001b[32m3.3999318915789445e-9\n",
      "\u001b[0m    Solve Status: \u001b[1m\u001b[32mSOLVE_SUCCEEDED\n"
     ]
    }
   ],
   "source": [
    "using Altro\n",
    "opts = SolverOptions(\n",
    "    cost_tolerance_intermediate=1e-2,\n",
    "    penalty_scaling=10.,\n",
    "    penalty_initial=1.0\n",
    ")\n",
    "\n",
    "altro = ALTROSolver(prob, opts)\n",
    "set_options!(altro, show_summary=true)\n",
    "solve!(altro);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can retrive some basic information about the solve for post-analysis,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "max_violation: 3.3999318915789445e-9\n",
      "cost:          1.552558743680986\n",
      "iterations:    40\n"
     ]
    }
   ],
   "source": [
    "println(\"max_violation: \", max_violation(altro))\n",
    "println(\"cost:          \", cost(altro))\n",
    "println(\"iterations:    \", iterations(altro));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and extract the solution from either `prob` or `altro` using `states` and `controls`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100-element Array{SArray{Tuple{1},Float64,1,1},1}:\n",
       " [-0.053404899938501706]\n",
       " [0.6325057323965376]\n",
       " [1.2322336398580338]\n",
       " [1.7121574166104716]\n",
       " [2.0447229731492684]\n",
       " [2.2096269379689777]\n",
       " [2.195634660099169]\n",
       " [2.0032216281222563]\n",
       " [1.6467798134436662]\n",
       " [1.154168025996269]\n",
       " [0.5622664906350295]\n",
       " [-0.09019963907039578]\n",
       " [-0.7690009035661038]\n",
       " ⋮\n",
       " [2.3242657214734757]\n",
       " [2.10770503727337]\n",
       " [1.8037656724994713]\n",
       " [1.455637884059381]\n",
       " [1.0817821499545335]\n",
       " [0.6836886129973376]\n",
       " [0.2516431323969652]\n",
       " [-0.23177751883430017]\n",
       " [-0.7901242106125287]\n",
       " [-1.4532878108851093]\n",
       " [-2.2592486680672152]\n",
       " [-3.000000000027063]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Extract the solution\n",
    "X = states(altro)\n",
    "U = controls(altro)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These functions return vectors of vectors. If you want to convert them to 2D arrays, simply use `hcat`. If `n`, `m`, or `N` are large, it's usually a good idea to convert it to a standard Julia `Vector` before concatenation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4×101 Array{Float64,2}:\n",
       " -3.05151e-11  -6.67562e-5    0.000591126  …  -0.00375642   8.87524e-11\n",
       "  8.12449e-11   0.000133512  -0.00118987       3.13402      3.14159\n",
       " -5.4539e-11   -0.00266588    0.0289336        0.150251    -2.19716e-10\n",
       " -2.84508e-11   0.0052881    -0.0576532        0.303027     7.06432e-10"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hcat(Vector.(X)...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solving with iLQR\n",
    "If we want to solve the unconstrained problem with iLQR, we simply reset our initial guess and create a new solver:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32;1m\n",
      "SOLVE COMPLETED\n",
      "\u001b[0m solved using the \u001b[0m\u001b[36;1miLQR\u001b[0m Solver,\n",
      " part of the Altro.jl package developed by the REx Lab at Stanford and Carnegie Mellon Universities\n",
      "\u001b[34;1m\n",
      "  Solve Statistics\n",
      "\u001b[0m    Total Iterations: 84\n",
      "\u001b[0m    Solve Time: 37.349365 (ms)\n",
      "\u001b[34;1m\n",
      "  Covergence\n",
      "\u001b[0m    Terminal Cost: 1.4497436179031664\n",
      "\u001b[0m    Terminal dJ: \u001b[32m6.889787558717053e-5\n",
      "\u001b[0m    Terminal gradient: \u001b[32m0.038402688096996665\n",
      "\u001b[0m    Solve Status: \u001b[1m\u001b[32mSOLVE_SUCCEEDED\n"
     ]
    }
   ],
   "source": [
    "ilqr = Altro.iLQRSolver(prob, opts)\n",
    "initial_controls!(ilqr, U0)\n",
    "solve!(ilqr);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solving with Ipopt\n",
    "TrajectoryOptimization provides a way to represent a trajectory optimization as a generic nonlinear program (NLP) and solve using solvers that implement the NLP interface in [`MathOptInterface`](https://github.com/jump-dev/MathOptInterface.jl), such as Ipopt (via [`Ipopt.jl`](https://github.com/jump-dev/Ipopt.jl)). \n",
    "\n",
    "Before converting to an NLP, we need to explicitly add the dynamics and initial condition as constraints to the optimization problem, which we do via `add_dynamics_constraints!`.\n",
    "\n",
    "We then convert our `Problem` to a `TrajOptNLP`, converting any applicable constraints to simple bounds on the decision variables (e.g. bound constraints and goal constraints), and specifying that the the constraint Jacobian uses a vector of non-zero values rather than a generic sparse matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "ename": "MethodError",
     "evalue": "MethodError: no method matching num_vars(::Int64, ::Int64, ::Int64, ::Bool)",
     "output_type": "error",
     "traceback": [
      "MethodError: no method matching num_vars(::Int64, ::Int64, ::Int64, ::Bool)",
      "",
      "Stacktrace:",
      " [1] TrajectoryOptimization.JacobianStructure(::ConstraintList, ::Symbol) at /home/bjack205/.julia/dev/TrajectoryOptimization/src/constraint_list.jl:265",
      " [2] TrajectoryOptimization.JacobianStructure(::ConstraintList) at /home/bjack205/.julia/dev/TrajectoryOptimization/src/constraint_list.jl:259",
      " [3] TrajOptNLP(::Problem{RK3,Float64}; remove_bounds::Bool, jac_type::Symbol) at /home/bjack205/.julia/dev/TrajectoryOptimization/src/nlp.jl:471",
      " [4] top-level scope at In[24]:14"
     ]
    }
   ],
   "source": [
    "using Ipopt\n",
    "using MathOptInterface\n",
    "const MOI = MathOptInterface\n",
    "\n",
    "# Copy problem to avoid modifying the original problem\n",
    "prob_nlp = copy(prob)\n",
    "\n",
    "# Add the dynamics and initial conditions as explicit constraints\n",
    "TrajectoryOptimization.add_dynamics_constraints!(prob_nlp)\n",
    "\n",
    "# Reset our initial guess\n",
    "initial_controls!(prob_nlp, U0)\n",
    "rollout!(prob_nlp)\n",
    "\n",
    "# Create the NLP\n",
    "nlp = TrajOptNLP(prob_nlp, remove_bounds=true, jac_type=:vector);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now solve with Ipopt by creating the `MathOptInterface.Optimizer` and setting up the problem in the `MathOptInterface` format using `TrajectoryOptimization.build_MOI!` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "******************************************************************************\n",
      "This program contains Ipopt, a library for large-scale nonlinear optimization.\n",
      " Ipopt is released as open source code under the Eclipse Public License (EPL).\n",
      "         For more information visit http://projects.coin-or.org/Ipopt\n",
      "******************************************************************************\n",
      "\n",
      "This is Ipopt version 3.13.2, running with linear solver mumps.\n",
      "NOTE: Other linear solvers might be more efficient (see Ipopt documentation).\n",
      "\n",
      "Number of nonzeros in equality constraint Jacobian...:     3568\n",
      "Number of nonzeros in inequality constraint Jacobian.:        0\n",
      "Number of nonzeros in Lagrangian Hessian.............:        0\n",
      "\n",
      "Total number of variables............................:      496\n",
      "                     variables with only lower bounds:        0\n",
      "                variables with lower and upper bounds:      100\n",
      "                     variables with only upper bounds:        0\n",
      "Total number of equality constraints.................:      400\n",
      "Total number of inequality constraints...............:        0\n",
      "        inequality constraints with only lower bounds:        0\n",
      "   inequality constraints with lower and upper bounds:        0\n",
      "        inequality constraints with only upper bounds:        0\n",
      "\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "   0  2.4696994e-01 3.14e+00 1.72e-05   0.0 0.00e+00    -  0.00e+00 0.00e+00   0\n",
      "   1  3.4719403e-01 2.53e+00 4.76e+00  -5.5 1.29e+01    -  2.76e-01 1.93e-01h  2\n",
      "   2  4.7772100e-01 2.22e+00 1.11e+01  -0.9 1.04e+01    -  4.42e-01 1.25e-01h  2\n",
      "   3  5.4482539e-01 2.08e+00 1.42e+01  -0.5 9.47e+00    -  5.29e-01 6.29e-02h  3\n",
      "   4  6.3290540e-01 1.91e+00 1.96e+01  -0.4 9.30e+00    -  7.08e-01 8.10e-02h  3\n",
      "   5  6.7395586e-01 1.83e+00 2.24e+01  -0.3 9.55e+00    -  7.54e-01 4.01e-02h  4\n",
      "   6  7.1665475e-01 1.75e+00 2.54e+01  -0.1 1.00e+01    -  7.24e-01 4.27e-02h  4\n",
      "   7  7.5767385e-01 1.68e+00 3.02e+01   0.1 1.01e+01    -  6.79e-01 4.17e-02h  4\n",
      "   8  7.9985755e-01 1.61e+00 3.58e+01   0.0 1.03e+01    -  7.37e-01 4.29e-02h  4\n",
      "   9  8.4641435e-01 1.53e+00 4.23e+01   0.1 1.04e+01    -  7.60e-01 4.73e-02h  4\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "  10  8.6842335e-01 1.49e+00 5.19e+01   0.4 1.08e+01    -  6.86e-01 2.45e-02h  5\n",
      "  11  8.9189767e-01 1.45e+00 6.23e+01   0.5 1.10e+01    -  6.51e-01 2.62e-02h  5\n",
      "  12  9.1660242e-01 1.41e+00 7.29e+01   0.6 1.12e+01    -  5.64e-01 2.75e-02h  5\n",
      "  13  9.4227933e-01 1.37e+00 8.80e+01   0.7 1.11e+01    -  6.10e-01 2.82e-02h  5\n",
      "  14  9.6757978e-01 1.34e+00 1.12e+02   0.7 1.10e+01    -  7.11e-01 2.72e-02h  5\n",
      "  15  9.9162633e-01 1.30e+00 1.35e+02   0.7 1.10e+01    -  5.55e-01 2.59e-02h  5\n",
      "  16  1.0153393e+00 1.27e+00 1.71e+02   0.8 1.07e+01    -  6.18e-01 2.54e-02h  5\n",
      "  17  1.0368639e+00 1.24e+00 2.10e+02   0.6 1.06e+01    -  5.46e-01 2.30e-02h  5\n",
      "  18  1.0578375e+00 1.21e+00 2.57e+02   0.7 1.04e+01    -  5.29e-01 2.32e-02h  5\n",
      "  19  1.0774436e+00 1.18e+00 3.14e+02   0.9 1.17e+01    -  4.97e-01 2.22e-02h  5\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "  20  1.1115488e+00 1.14e+00 3.50e+02   1.0 1.19e+01    -  3.06e-01 3.63e-02h  4\n",
      "  21  1.1391461e+00 1.11e+00 3.77e+02   0.8 1.47e+01    -  2.18e-01 2.83e-02h  4\n",
      "  22  1.1637858e+00 1.08e+00 4.09e+02   0.8 1.78e+01    -  2.18e-01 2.54e-02h  4\n",
      "  23  1.1879277e+00 1.05e+00 4.55e+02   1.2 1.86e+01    -  2.79e-01 2.74e-02h  4\n",
      "  24  1.2334740e+00 1.01e+00 4.98e+02   1.1 2.28e+01    -  2.04e-01 4.33e-02h  3\n",
      "  25  1.2750733e+00 9.65e-01 6.45e+02   1.4 1.83e+01    -  1.98e-01 4.03e-02h  3\n",
      "  26  1.3022401e+00 9.30e-01 7.89e+02   1.6 2.46e+01    -  1.67e-01 3.68e-02h  3\n",
      "  27  1.3213479e+00 8.99e-01 9.20e+02   1.7 3.16e+01    -  3.05e-01 3.34e-02h  3\n",
      "  28  1.3584968e+00 8.58e-01 1.13e+03   1.2 2.24e+01    -  1.42e-01 4.56e-02h  3\n",
      "  29  1.3828911e+00 8.29e-01 1.31e+03   1.6 2.82e+01    -  4.28e-01 3.41e-02h  3\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "  30  1.3971535e+00 7.97e-01 1.65e+03   2.0 3.04e+01    -  4.14e-01 3.80e-02h  3\n",
      "  31  1.4209984e+00 7.62e-01 2.06e+03   1.9 2.19e+01    -  2.17e-01 4.43e-02h  3\n",
      "  32  1.4342773e+00 7.15e-01 2.66e+03   2.0 4.27e+01    -  3.46e-01 6.12e-02h  2\n",
      "  33  1.4416180e+00 6.96e-01 2.96e+03   2.1 1.76e+01    -  4.94e-01 2.66e-02h  4\n",
      "  34  1.4553510e+00 6.72e-01 3.34e+03   2.1 1.58e+01    -  6.92e-01 3.51e-02h  4\n",
      "  35  1.4674354e+00 6.47e-01 3.84e+03   2.2 1.55e+01    -  5.54e-01 3.67e-02h  4\n",
      "  36  1.4754523e+00 6.25e-01 4.31e+03   1.9 2.04e+01    -  3.19e-01 3.41e-02h  4\n",
      "  37  1.4890706e+00 5.97e-01 5.17e+03   2.0 1.39e+01    -  1.74e-01 4.54e-02h  4\n",
      "  38  1.5036117e+00 5.64e-01 6.66e+03   2.2 2.05e+01    -  5.34e-01 5.46e-02h  3\n",
      "  39  1.5307477e+00 5.19e-01 8.98e+03   2.5 1.37e+01    -  6.18e-01 8.04e-02h  3\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "  40  1.5500449e+00 4.89e-01 1.11e+04   2.7 6.24e+00    -  7.46e-01 5.63e-02h  4\n",
      "  41  1.5532055e+00 4.66e-01 1.28e+04   2.4 2.96e+01    -  3.64e-01 4.73e-02h  3\n",
      "  42  1.5737279e+00 4.42e-01 1.50e+04   2.2 1.03e+01    -  2.55e-01 5.27e-02h  4\n",
      "  43  1.5954913e+00 4.11e-01 1.91e+04   2.5 1.38e+01    -  8.68e-01 7.02e-02h  3\n",
      "  44  1.6085518e+00 3.83e-01 2.42e+04   2.7 1.83e+01    -  7.04e-01 6.85e-02h  3\n",
      "  45  1.6289623e+00 3.44e-01 3.28e+04   3.1 1.09e+01    -  5.55e-01 1.01e-01h  3\n",
      "  46  1.6378255e+00 3.16e-01 3.44e+04   3.3 6.46e+01    -  4.00e-01 8.07e-02h  1\n",
      "  47  1.6420093e+00 3.01e-01 4.03e+04   3.2 1.23e+01    -  6.56e-01 4.66e-02f  4\n",
      "  48  1.6613403e+00 2.80e-01 4.86e+04   3.2 5.49e+00    -  8.76e-01 7.20e-02h  4\n",
      "  49  1.6816237e+00 2.56e-01 6.32e+04   3.2 1.31e+01    -  7.77e-01 8.44e-02h  3\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "  50  1.7038297e+00 2.33e-01 7.94e+04   2.6 1.54e+01    -  4.08e-01 9.07e-02h  3\n",
      "  51  1.7042157e+00 2.16e-01 9.22e+04   3.4 1.48e+01    -  8.50e-01 7.09e-02h  3\n",
      "  52  1.7571521e+00 1.78e-01 1.10e+05   3.4 1.25e+01    -  8.99e-01 1.79e-01h  2\n",
      "  53  1.7779503e+00 1.63e-01 1.28e+05   3.4 5.83e+00    -  7.82e-01 8.44e-02h  4\n",
      "  54  1.8027567e+00 1.46e-01 1.55e+05   2.8 1.20e+01    -  3.25e-01 1.03e-01h  3\n",
      "  55  1.8105657e+00 1.27e-01 1.80e+05   3.3 4.32e+01    -  4.35e-01 1.27e-01h  1\n",
      "  56  1.8112687e+00 1.15e-01 2.03e+05   3.3 1.31e+01    -  9.68e-01 9.32e-02f  3\n",
      "  57  1.8400203e+00 9.57e-02 2.02e+05   3.1 4.23e+00    -  9.91e-01 1.71e-01h  3\n",
      "  58  2.0118042e+00 4.77e-02 1.50e+05   2.5 9.75e+00    -  5.14e-01 5.02e-01h  1\n",
      "  59  2.0788675e+00 2.12e-02 7.51e+04   1.9 7.49e+00    -  2.73e-01 8.29e-01H  1\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "  60  2.0765787e+00 1.39e-02 2.72e+04   2.2 3.64e+00    -  9.30e-01 7.09e-01f  1\n",
      "  61  2.0837858e+00 5.06e-04 4.14e+02  -3.7 2.08e-01    -  8.46e-01 1.00e+00h  1\n",
      "  62  2.0775977e+00 1.25e-05 1.11e+01  -0.5 2.03e-01    -  9.80e-01 1.00e+00f  1\n",
      "  63  2.0723815e+00 4.13e-06 4.41e-01  -2.1 7.35e-02    -  9.62e-01 1.00e+00f  1\n",
      "  64  1.9852836e+00 1.06e-03 3.20e-01  -3.3 1.02e+00    -  3.48e-01 1.00e+00f  1\n",
      "  65  1.8763537e+00 7.28e-03 3.05e-01  -2.7 3.21e+00    -  5.31e-01 1.00e+00f  1\n",
      "  66  1.7975420e+00 2.47e-03 3.91e-01  -2.6 1.46e+00    -  6.00e-01 1.00e+00h  1\n",
      "  67  1.7437302e+00 1.13e-02 4.17e-01  -2.4 2.78e+00    -  9.06e-01 8.68e-01h  1\n",
      "  68  1.7043553e+00 5.14e-03 4.00e-01  -2.7 1.95e+00    -  6.19e-01 1.00e+00h  1\n",
      "  69  1.6469904e+00 1.26e-01 4.83e-01  -2.6 1.53e+00    -  9.32e-01 1.00e+00h  1\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "  70  1.6817345e+00 1.25e-01 7.12e-01  -1.9 7.10e+00    -  8.49e-01 2.91e-01h  1\n",
      "  71  1.6645347e+00 5.94e-02 5.47e-01  -2.1 7.10e+00    -  9.97e-01 6.14e-01h  1\n",
      "  72  1.6176994e+00 5.44e-03 2.19e-01  -2.2 1.46e+00    -  9.85e-01 1.00e+00h  1\n",
      "  73  1.5942720e+00 1.39e-02 1.44e-01  -2.6 1.96e+00    -  9.91e-01 1.00e+00h  1\n",
      "  74  1.5782236e+00 1.14e-02 2.31e-01  -2.7 4.50e+00    -  7.65e-01 1.00e+00H  1\n",
      "  75  1.5723246e+00 2.94e-03 1.65e-01  -2.4 3.81e+00    -  1.00e+00 1.00e+00h  1\n",
      "  76  1.5542846e+00 2.00e-03 1.08e-01  -2.9 3.29e-01    -  9.92e-01 1.00e+00h  1\n",
      "  77  1.5376735e+00 2.64e-02 1.08e-01  -3.3 9.06e-01    -  1.00e+00 1.00e+00h  1\n",
      "  78  1.5258039e+00 2.10e-02 2.35e-01  -3.0 1.90e+00    -  6.16e-01 1.00e+00H  1\n",
      "  79  1.5233440e+00 1.18e-02 2.95e-01  -3.3 2.30e+00    -  1.00e+00 7.56e-01h  1\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "  80  1.5528946e+00 3.23e-03 7.50e-02  -2.8 1.75e+00    -  1.00e+00 1.00e+00H  1\n",
      "  81  1.5174427e+00 4.23e-04 5.85e-02  -3.0 1.16e+00    -  9.99e-01 1.00e+00h  1\n",
      "  82  1.5075281e+00 1.69e-04 3.22e-02  -3.7 4.02e-01    -  9.96e-01 1.00e+00h  1\n",
      "  83  1.5025143e+00 1.06e-03 3.34e-02  -4.2 3.43e-01    -  9.98e-01 1.00e+00h  1\n",
      "  84  1.4999054e+00 2.26e-04 4.27e-02  -4.3 6.08e-01    -  1.00e+00 1.00e+00H  1\n",
      "  85  1.4987531e+00 4.45e-03 7.16e-02  -5.2 3.33e-01    -  1.00e+00 1.00e+00h  1\n",
      "  86  1.5005847e+00 5.76e-04 3.23e-02  -4.2 5.44e-01    -  1.00e+00 1.00e+00h  1\n",
      "  87  1.4974698e+00 1.18e-03 4.17e-02  -4.3 2.58e-01    -  1.00e+00 1.00e+00h  1\n",
      "  88  1.4976326e+00 8.73e-04 2.22e-02  -4.4 1.27e-01    -  9.99e-01 1.00e+00h  1\n",
      "  89  1.4976153e+00 9.65e-04 2.96e-02  -4.1 2.65e+00    -  1.00e+00 2.83e-02h  5\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "  90  1.4978167e+00 4.93e-03 3.87e-02  -3.9 5.61e-01    -  1.00e+00 1.00e+00H  1\n",
      "  91  1.4982595e+00 1.52e-04 3.39e-02  -4.0 4.17e-01    -  1.00e+00 1.00e+00h  1\n",
      "  92  1.4965754e+00 3.98e-05 1.54e-02  -4.6 7.66e-02    -  1.00e+00 1.00e+00h  1\n",
      "  93  1.4961549e+00 7.11e-05 1.14e-02  -5.2 5.95e-02    -  1.00e+00 1.00e+00h  1\n",
      "  94  1.4959387e+00 3.90e-05 2.13e-02  -6.0 1.09e-01    -  9.95e-01 1.00e+00H  1\n",
      "  95  1.4972658e+00 2.16e-06 2.46e-02  -5.3 5.14e-01    -  8.81e-01 1.00e+00H  1\n",
      "  96  1.4966470e+00 6.17e-04 8.06e-03  -5.9 4.35e-01    -  1.00e+00 1.00e+00h  1\n",
      "  97  1.4959122e+00 3.80e-05 5.06e-03  -6.0 4.42e-02    -  1.00e+00 1.00e+00h  1\n",
      "  98  1.4958923e+00 2.18e-05 4.25e-03  -6.6 3.93e-02    -  1.00e+00 1.00e+00h  1\n",
      "  99  1.4958936e+00 1.53e-05 3.33e-03  -7.9 2.97e-02    -  1.00e+00 1.00e+00h  1\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      " 100  1.4958787e+00 7.18e-06 7.18e-04  -8.9 2.45e-02    -  1.00e+00 1.00e+00h  1\n",
      " 101  1.4958743e+00 6.46e-07 8.40e-04  -9.7 1.15e-02    -  1.00e+00 1.00e+00h  1\n",
      " 102  1.4958740e+00 2.23e-07 2.66e-05 -10.9 5.33e-03    -  1.00e+00 1.00e+00h  1\n",
      " 103  1.4958739e+00 3.94e-10 8.29e-07 -11.0 3.94e-04    -  1.00e+00 1.00e+00h  1\n",
      " 104  1.4958739e+00 3.90e-11 1.01e-07 -11.0 1.75e-04    -  1.00e+00 1.00e+00h  1\n",
      " 105  1.4958739e+00 7.54e-12 7.96e-08 -11.0 1.98e-05    -  1.00e+00 1.00e+00h  1\n",
      " 106  1.4958739e+00 8.87e-13 3.30e-08 -11.0 9.63e-06    -  1.00e+00 1.00e+00h  1\n",
      " 107  1.4958739e+00 1.78e-15 1.66e-07 -11.0 4.22e-05    -  1.00e+00 1.00e+00H  1\n",
      " 108  1.4958739e+00 1.78e-15 5.94e-07 -11.0 1.21e-04    -  1.00e+00 1.00e+00H  1\n",
      " 109  1.4958739e+00 2.75e-14 1.23e-07 -11.0 1.12e-04    -  1.00e+00 1.00e+00H  1\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      " 110  1.4958739e+00 1.78e-15 5.39e-07 -11.0 1.60e-04    -  1.00e+00 1.00e+00H  1\n",
      " 111  1.4958739e+00 2.66e-14 1.48e-07 -11.0 1.42e-04    -  1.00e+00 1.00e+00H  1\n",
      " 112  1.4958739e+00 1.78e-15 5.06e-07 -11.0 1.21e-04    -  1.00e+00 1.00e+00H  1\n",
      " 113  1.4958739e+00 5.68e-14 1.72e-07 -11.0 8.82e-05    -  1.00e+00 1.00e+00H  1\n",
      " 114  1.4958739e+00 1.66e-12 9.03e-08 -11.0 2.23e-04    -  1.00e+00 1.25e-01h  4\n",
      " 115  1.4958739e+00 2.62e-12 8.45e-08 -11.0 2.99e-05    -  1.00e+00 1.00e+00h  1\n",
      " 116  1.4958739e+00 3.69e-12 2.42e-08 -11.0 1.02e-05    -  1.00e+00 1.00e+00h  1\n",
      " 117  1.4958739e+00 1.24e-13 5.35e-09 -11.0 2.60e-06    -  1.00e+00 1.00e+00h  1\n",
      "\n",
      "Number of Iterations....: 117\n",
      "\n",
      "                                   (scaled)                 (unscaled)\n",
      "Objective...............:   1.4958739082433778e+00    1.4958739082433778e+00\n",
      "Dual infeasibility......:   5.3454160109170440e-09    5.3454160109170440e-09\n",
      "Constraint violation....:   1.2434497875801753e-13    1.2434497875801753e-13\n",
      "Complementarity.........:   1.0000003923437651e-11    1.0000003923437651e-11\n",
      "Overall NLP error.......:   5.3454160109170440e-09    5.3454160109170440e-09\n",
      "\n",
      "\n",
      "Number of objective function evaluations             = 425\n",
      "Number of objective gradient evaluations             = 118\n",
      "Number of equality constraint evaluations            = 425\n",
      "Number of inequality constraint evaluations          = 0\n",
      "Number of equality constraint Jacobian evaluations   = 118\n",
      "Number of inequality constraint Jacobian evaluations = 0\n",
      "Number of Lagrangian Hessian evaluations             = 0\n",
      "Total CPU secs in IPOPT (w/o function evaluations)   =      2.210\n",
      "Total CPU secs in NLP function evaluations           =      0.402\n",
      "\n",
      "EXIT: Optimal Solution Found.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "LOCALLY_SOLVED::TerminationStatusCode = 4"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optimizer = Ipopt.Optimizer()\n",
    "TrajectoryOptimization.build_MOI!(nlp, optimizer)\n",
    "MOI.optimize!(optimizer)\n",
    "MOI.get(optimizer, MOI.TerminationStatus())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the solve updates the data structures in `nlp` during the course of the solve, we can use either `optimizer` or `nlp` to get information after the solve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "max_violation: 1.2434497875801753e-13\n",
      "cost:          1.4958739082433778\n"
     ]
    }
   ],
   "source": [
    "println(\"max_violation: \", max_violation(nlp))\n",
    "println(\"cost:          \", cost(nlp));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparing Solve Times\n",
    "Here we demonstrate the compuational advantage of ALTRO over Ipopt. We demonstrate how to use `BenchmarkTools` to time a solve using `MathOptInterface` by resetting the initial guess between solves. As shown, ALTRO is about 100x faster than Ipopt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BenchmarkTools.Trial: \n",
       "  memory estimate:  1.52 MiB\n",
       "  allocs estimate:  1581\n",
       "  --------------\n",
       "  minimum time:     6.547 ms (0.00% GC)\n",
       "  median time:      7.132 ms (0.00% GC)\n",
       "  mean time:        7.095 ms (1.28% GC)\n",
       "  maximum time:     7.588 ms (4.16% GC)\n",
       "  --------------\n",
       "  samples:          10\n",
       "  evals/sample:     10"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using BenchmarkTools\n",
    "\n",
    "# Reset initial guess and then benchmark ALTRO solver\n",
    "initial_controls!(altro, U0)\n",
    "b_altro = benchmark_solve!(altro)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BenchmarkTools.Trial: \n",
       "  memory estimate:  7.79 MiB\n",
       "  allocs estimate:  22123\n",
       "  --------------\n",
       "  minimum time:     619.169 ms (0.00% GC)\n",
       "  median time:      693.174 ms (0.00% GC)\n",
       "  mean time:        752.114 ms (0.08% GC)\n",
       "  maximum time:     1.220 s (0.00% GC)\n",
       "  --------------\n",
       "  samples:          7\n",
       "  evals/sample:     1"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Reset initial guess and benchmark Ipopt solver\n",
    "initial_controls!(prob_nlp, U0)\n",
    "rollout!(prob_nlp)\n",
    "nlp = TrajOptNLP(prob_nlp, remove_bounds=true, jac_type=:vector)\n",
    "Z0 = copy(TrajectoryOptimization.get_trajectory(nlp).Z)\n",
    "\n",
    "optimizer = Ipopt.Optimizer(print_level=0)\n",
    "TrajectoryOptimization.build_MOI!(nlp, optimizer)\n",
    "Z = MOI.VariableIndex.(1:length(Z0))\n",
    "b_ipopt = @benchmark begin\n",
    "    MOI.optimize!($optimizer)\n",
    "    MOI.set($optimizer, MOI.VariablePrimalStart(), $Z, $Z0)\n",
    "    MOI.get($optimizer, MOI.TerminationStatus())\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cost(nlp) = 1.4958739082433778\n",
      "cost(altro) = 1.5525585360226632\n",
      "max_violation(nlp) = 1.2434497875801753e-13\n",
      "max_violation(altro) = 3.4246698810136422e-9\n",
      "Speed improvement: 97.0x\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BenchmarkTools.TrialJudgement: \n",
       "  time:   -98.97% => \u001b[32mimprovement\u001b[39m (5.00% tolerance)\n",
       "  memory: -80.46% => \u001b[32mimprovement\u001b[39m (1.00% tolerance)\n"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Statistics\n",
    "@show cost(nlp)\n",
    "@show cost(altro)\n",
    "@show max_violation(nlp)\n",
    "@show max_violation(altro)\n",
    "jdg = judge(median(b_altro), median(b_ipopt))\n",
    "println(\"Speed improvement: \", round(1/ratio(jdg).time), \"x\")\n",
    "jdg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing the Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "┌ Info: MeshCat server started. You can open the visualizer by visiting the following URL in your browser:\n",
      "│ http://localhost:8702\n",
      "└ @ MeshCat /home/bjack205/.julia/packages/MeshCat/ECbzr/src/visualizer.jl:73\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "    <div style=\"height: 500px; width: 100%; overflow-x: auto; overflow-y: hidden; resize: both\">\n",
       "    <iframe src=\"http://localhost:8702\" style=\"width: 100%; height: 100%; border: none\"></iframe>\n",
       "    </div>\n"
      ],
      "text/plain": [
       "MeshCat.DisplayedVisualizer(MeshCat.CoreVisualizer(MeshCat.SceneTrees.SceneNode(nothing, nothing, Dict{String,Array{UInt8,1}}(), nothing, Dict{String,MeshCat.SceneTrees.SceneNode}()), Set(Any[]), ip\"127.0.0.1\", 8702))"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using TrajOptPlots\n",
    "using MeshCat\n",
    "using Plots\n",
    "\n",
    "vis = Visualizer()\n",
    "render(vis)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the default geometries for the robot, as defined in `TrajOptPlots`, using the `set_mesh!` command"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MeshCat Visualizer with path /meshcat/robot/cart/pole at http://localhost:8702"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "TrajOptPlots.set_mesh!(vis, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the results of a solver by passing the solver to `visualize!`. In practice, all a solver needs to implement for this method is a `get_model` method that returns an `AbstractModel` and a `get_trajectory` method that returns an `AbstractTrajectory`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize!(vis, altro);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize!(vis, model, TrajectoryOptimization.get_trajectory(nlp));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also compare both solutions by passing in both solvers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize!(vis, altro, nlp);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or the model and different trajectories:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize!(vis, model, get_trajectory(altro), get_trajectory(nlp));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After this, we'll have extra geometries floating around, which can be deleting using:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "delete!(vis[\"robot_copies\"]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tip: Defining the visualazation method\n",
    "Visualization of the model is defined by the `visualize!(vis, model::MyModel, x::AbstractVector)` command:\n",
    " \n",
    " ```julia\n",
    " function TrajOptPlots.visualize!(vis, model::RobotZoo.Cartpole, x::AbstractVector)\n",
    "    y = x[1]\n",
    "    θ = x[2]\n",
    "    q = expm((pi-θ) * @SVector [1,0,0])\n",
    "    settransform!(vis[\"robot\",\"cart\"], Translation(0,-y,0))\n",
    "    settransform!(vis[\"robot\",\"cart\",\"pole\",\"geom\"], LinearMap(UnitQuaternion(q)))\n",
    "end\n",
    "\n",
    " ```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.4.2",
   "language": "julia",
   "name": "julia-1.4"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.4.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
